
library(caret)
library(rpart)

source("helpers.R")

# loading data
load("../features/features_combined.Rdata")
source("../load_data.R")

train <- merge(train_features, trainBlockB, by = "STUDENTID")
test <- test_features

# 设置训练目标
train$efficient <- as.factor(train$`mode(EfficientlyCompletedBlockB)`)
train$`mode(EfficientlyCompletedBlockB)`<-NULL
train$STUDENTID<-NULL
levels(train$efficient) <- make.names(levels(train$efficient))

train_control <- trainControl(method = "cv", number = 10, classProbs = TRUE, summaryFunction = twoClassSummary)

parameter_grid_gbm <- expand.grid(interaction.depth = c(1, 2), n.trees = seq(200, 1000, by = 50),
                                  shrinkage=c(0.03, 0.02, 0.01, 0.0075, 0.005),
                                  n.minobsinnode = c(30, 35, 40, 45, 50))
parameter_grid_RRF <- expand.grid(mtry = seq(2, 15, by = 2), coefImp = 0.0, coefReg = c(0.7, 0.85, 1.0))
parameter_grid_rf <- expand.grid(mtry = seq(2, 15, by = 2))
parameter_grid_regLogistic <- NULL
parameter_grid_dwdPoly <- expand.grid(lambda = c(0.005, 0.01), qval = seq(0.00, 0.5, by = 0.05),
                                      degree = 2, scale = 0.01)
parameter_grid_knn <- expand.grid(k = seq(10, 70, by = 5))
# parameter_grid_nb <- expand.grid(
#   laplace = c(0, 1),
#   usekernel = c(TRUE, FALSE),
#   adjust = seq(0, 2, by = 0.5)
# )
parameter_grid_pls <- expand.grid(ncomp = seq(1, 10))
parameter_grid_svmRadial <- expand.grid(C = seq(0.1, 1.0, by = 0.1), sigma = 2 ^ seq(-10, -5, length.out = 10))
parameter_grid_dwdLinear <- expand.grid(qval = seq(0.01, 1, by = 0.05), lambda = 10 ^ seq(-6, -2, by = 1))
parameter_grid_nnet <- expand.grid(size = seq(5, 9, by = 1), decay = seq(1, 5, length.out = 15))



# 模型配置
models <- list(
  gbm = list(method = "gbm", train_control = train_control, tuneGrid = parameter_grid_gbm),
  RRF = list(method = "RRF", train_control = train_control, tuneGrid = parameter_grid_RRF),
  rf = list(method = "rf", train_control = train_control, tuneGrid = parameter_grid_rf),
  regLogistic = list(method = "regLogistic", train_control = train_control, tuneGrid = NULL), # 逻辑回归通常不需要网格搜索
  dwdPoly = list(method = "dwdPoly", train_control = train_control, tuneGrid = parameter_grid_dwdPoly),
  knn = list(method = "knn", train_control = train_control, pre_process = c('center', 'scale'), tuneGrid = parameter_grid_knn),
  # nb = list(method = "naive_bayes", train_control = train_control)
  pls = list(method = "pls", train_control = train_control, tuneGrid = parameter_grid_pls),
  svmRadial = list(method = "svmRadial", train_control = train_control, tuneGrid = parameter_grid_svmRadial),
  dwdLinear = list(method = "dwdLinear", train_control = train_control, tuneGrid = parameter_grid_dwdLinear),
  nnet = list(method = "nnet", train_control = train_control, tuneGrid = parameter_grid_nnet)
)

# 更新函数以使用参数网格
train_and_evaluate_models <- function(data, models, outcome_var) {
  results <- list()
  
  for(model_name in names(models)) {
    cat("Training model:", model_name, "\n")
    model <- train(reformulate(names(data)[!names(data) %in% outcome_var], response = outcome_var),
                   data = data,
                   method = models[[model_name]]$method,
                   trControl = models[[model_name]]$train_control,
                   metric = "ROC",
                   preProcess = models[[model_name]]$pre_process,
                   tuneGrid = models[[model_name]]$tuneGrid)
    
    cat("Model trained. Making predictions...\n")
    predicted_classes <- predict(model, newdata = data, type = "raw")
    conf_matrix <- confusionMatrix(predicted_classes, data[[outcome_var]])
    
    cat("Model", model_name, "Confusion Matrix:\n")
    print(conf_matrix)
    
    results[[model_name]] <- list(model = model, confusionMatrix = conf_matrix)
  }
  
  return(results)
}

# 使用函数
results <- train_and_evaluate_models(train, models, "efficient")

print(results)
save(results,
     file = "model.Rdata")
